#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
replace_tif_values.py

功能：
- 读取 Excel 映射表，将 GeoTIFF 像元值按映射关系进行替换；
- 自动生成示例数据（一个小的示例 GeoTIFF 与一个 Excel 对应表），便于快速试跑；
- 保留原始栅格的空间参考、仿射变换、nodata 等元数据；

依赖：rasterio, numpy, pandas, openpyxl

用法示例：
1) 生成示例数据：
   python replace_tif_values.py --make-sample --outdir data

2) 使用 Excel 对应表替换：
   python replace_tif_values.py \
     --input data/sample_input.tif \
     --excel data/sample_mapping.xlsx \
     --output data/sample_output.tif

可选参数：
- --sheet 指定 Excel 的工作表名（默认 Sheet1）
- --from-col 映射“原值”列名（默认 from）
- --to-col 映射“新值”列名（默认 to）
- --apply-band 指定只处理的波段序号（从1开始）；不提供则处理全部波段
- --ignore-missing 在映射中找不到的值保持不变（默认行为）。如需严格模式，使用 --strict
- --strict 开启严格模式：如果出现 Excel 未覆盖的值则报错退出

注意：
- 本脚本默认采用 LZW 压缩输出；
- Excel 映射列会尝试转为整数；若包含小数将以浮点处理；
- 仅对数值型像元做映射，nodata 像元保持不变。
"""

import argparse
import os
from typing import Dict, Optional, Tuple, List

import numpy as np
import pandas as pd
import rasterio
from rasterio.transform import from_origin


def read_mapping_from_excel(excel_path: str, sheet: str, from_col: str, to_col: str) -> Tuple[Dict[float, float], bool]:
	"""读取 Excel 映射表，返回字典和是否为整数映射的标记。

	尝试以整数读取；若任一列含非整数，将以浮点处理。
	"""
	df = pd.read_excel(excel_path, sheet_name=sheet)
	if from_col not in df.columns or to_col not in df.columns:
		raise ValueError(f"Excel中找不到列: {from_col} 或 {to_col}. 实际列: {list(df.columns)}")

	# 判断是否可整体转换为整数
	def can_cast_to_int(series: pd.Series) -> bool:
		try:
			_ = series.dropna().astype(int)
			return True
		except Exception:
			return False

	from_series = df[from_col]
	to_series = df[to_col]

	all_int = can_cast_to_int(from_series) and can_cast_to_int(to_series)
	if all_int:
		from_vals = from_series.dropna().astype(int).tolist()
		to_vals = to_series.dropna().astype(int).tolist()
		mapping = {int(k): int(v) for k, v in zip(from_vals, to_vals)}
	else:
		from_vals = from_series.dropna().astype(float).tolist()
		to_vals = to_series.dropna().astype(float).tolist()
		mapping = {float(k): float(v) for k, v in zip(from_vals, to_vals)}

	return mapping, all_int


def _select_output_dtype(all_int: bool, to_values: List[float]) -> np.dtype:
	"""根据映射目标值选择合适的输出 dtype。"""
	if not all_int:
		return np.dtype("float32")
	min_v = int(min(to_values)) if to_values else 0
	max_v = int(max(to_values)) if to_values else 0
	if min_v >= 0:
		# 无符号
		if max_v <= 255:
			return np.dtype("uint8")
		elif max_v <= 65535:
			return np.dtype("uint16")
		elif max_v <= 4294967295:
			return np.dtype("uint32")
		else:
			return np.dtype("float32")
	else:
		# 有符号
		if -32768 <= min_v and max_v <= 32767:
			return np.dtype("int16")
		elif -2147483648 <= min_v and max_v <= 2147483647:
			return np.dtype("int32")
		else:
			return np.dtype("float32")


def apply_mapping_to_array(array: np.ndarray, mapping: Dict[float, float], nodata: Optional[float], strict: bool, out_dtype: Optional[np.dtype] = None) -> np.ndarray:
	"""对二维或三维数组应用映射。

	- 支持 (bands, height, width) 或 (height, width) 形状；
	- nodata 像元保持不变；
	- 严格模式会检查是否存在未覆盖的值。
	"""
	if array.ndim == 2:
		arrays = array[np.newaxis, ...]
		squeeze_back = True
	elif array.ndim == 3:
		arrays = array
		squeeze_back = False
	else:
		raise ValueError("仅支持2D或3D数组")

	result = arrays.copy()

	# 构造掩膜：有效像元（非 nodata）
	if nodata is not None:
		valid_mask = ~np.isclose(arrays, nodata)
	else:
		valid_mask = np.ones_like(arrays, dtype=bool)

	# 在替换前转换 dtype 以避免溢出
	if out_dtype is not None and result.dtype != out_dtype:
		result = result.astype(out_dtype, copy=False)

	if strict:
		# 找出所有有效像元的唯一值
		unique_vals = np.unique(arrays[valid_mask])
		missing = [v for v in unique_vals if v not in mapping]
		if len(missing) > 0:
			raise ValueError(f"严格模式下发现未在Excel映射中的值: {missing[:20]}... 共{len(missing)}个")

	# 逐键替换（对稀疏映射较稳定）
	for src_val, dst_val in mapping.items():
		mask = valid_mask & np.isclose(arrays, src_val)
		result[mask] = dst_val

	if squeeze_back:
		return result[0]
	return result


def write_raster(output_path: str, array: np.ndarray, meta: dict) -> None:
	"""写出 GeoTIFF，保留原始元数据并使用 LZW 压缩。"""
	out_meta = meta.copy()
	out_meta.update({
		"compress": "lzw"
	})
	# rasterio 写入需要 (bands, height, width)
	if array.ndim == 2:
		write_arr = array[np.newaxis, ...]
	else:
		write_arr = array
	# 同步 dtype
	out_meta.update({
		"count": write_arr.shape[0],
		"dtype": np.dtype(write_arr.dtype).name
	})
	with rasterio.open(output_path, "w", **out_meta) as dst:
		for bidx in range(write_arr.shape[0]):
			dst.write(write_arr[bidx, ...], bidx + 1)


def make_sample_data(outdir: str) -> Tuple[str, str]:
	"""生成示例 GeoTIFF 与 Excel 映射表，返回两个文件路径。"""
	os.makedirs(outdir, exist_ok=True)

	# 生成 100x100 的整数分类栅格（值域 0..5，其中 0 作为 nodata）
	height, width = 100, 100
	transform = from_origin(120.0, 40.0, 0.01, 0.01)  # 假定经纬度仿射，仅示例
	nodata_val = 0
	rng = np.random.default_rng(42)
	data = rng.integers(low=0, high=6, size=(height, width), dtype=np.int16)

	# 写出示例 TIF
	crs = "EPSG:4326"
	tif_path = os.path.join(outdir, "sample_input.tif")
	meta = {
		"driver": "GTiff",
		"height": height,
		"width": width,
		"count": 1,
		"dtype": "int16",
		"crs": crs,
		"transform": transform,
		"nodata": nodata_val,
		"compress": "lzw",
	}
	with rasterio.open(tif_path, "w", **meta) as dst:
		dst.write(data, 1)

	# 生成 Excel 映射：1->10, 2->20, 3->30, 4->40, 5->50（0 作为 nodata 不映射）
	mapping_df = pd.DataFrame({
		"from": [1, 2, 3, 4, 5],
		"to":   [10, 20, 30, 40, 50],
	})
	excel_path = os.path.join(outdir, "sample_mapping.xlsx")
	with pd.ExcelWriter(excel_path, engine="openpyxl") as writer:
		mapping_df.to_excel(writer, sheet_name="Sheet1", index=False)

	return tif_path, excel_path



def process(input_tif: str, excel_path: str, output_tif: str, sheet: str, from_col: str, to_col: str, apply_band: Optional[int], ignore_missing: bool, strict: bool) -> None:
	mapping, all_int = read_mapping_from_excel(excel_path, sheet=sheet, from_col=from_col, to_col=to_col)
	# 确定输出 dtype
	to_values = list(mapping.values())
	out_dtype = _select_output_dtype(all_int=all_int, to_values=to_values)
	with rasterio.open(input_tif) as src:
		meta = src.meta.copy()
		nodata = src.nodata
		if apply_band is None:
			arr = src.read()  # (bands, H, W)
			result = apply_mapping_to_array(arr, mapping=mapping, nodata=nodata, strict=strict, out_dtype=out_dtype)
		else:
			if apply_band < 1 or apply_band > src.count:
				raise ValueError(f"apply-band 超出范围，应在 1..{src.count} 之间")
			arr = src.read(apply_band)  # (H, W)
			res_band = apply_mapping_to_array(arr, mapping=mapping, nodata=nodata, strict=strict, out_dtype=out_dtype)
			# 将结果放回对应波段，其余波段保持不变
			full = src.read()
			if full.dtype != out_dtype:
				full = full.astype(out_dtype)
			full[apply_band - 1, :, :] = res_band
			result = full

	write_raster(output_tif, result, meta)


def list_tif_files(input_dir: str) -> List[str]:
	"""递归列出目录下所有 .tif/.tiff 文件。"""
	result: List[str] = []
	for root, _dirs, files in os.walk(input_dir):
		for name in files:
			lower = name.lower()
			if lower.endswith(".tif") or lower.endswith(".tiff"):
				result.append(os.path.join(root, name))
	return result


def process_directory(input_dir: str, excel_path: str, output_dir: str, sheet: str, from_col: str, to_col: str, apply_band: Optional[int], ignore_missing: bool, strict: bool) -> List[str]:
	"""批量处理目录下所有 TIF 文件，返回成功输出的文件路径列表。"""
	os.makedirs(output_dir, exist_ok=True)
	files = list_tif_files(input_dir)
	if not files:
		print(f"目录内未找到 TIF 文件：{input_dir}")
		return []

	# 读取一次映射，避免重复读取 Excel
	mapping, all_int = read_mapping_from_excel(excel_path, sheet=sheet, from_col=from_col, to_col=to_col)
	to_values = list(mapping.values())
	out_dtype = _select_output_dtype(all_int=all_int, to_values=to_values)
	written: List[str] = []
	for in_path in files:
		basename = os.path.basename(in_path)
		out_path = os.path.join(output_dir, basename)
		with rasterio.open(in_path) as src:
			meta = src.meta.copy()
			nodata = src.nodata
			if apply_band is None:
				arr = src.read()
				res = apply_mapping_to_array(arr, mapping=mapping, nodata=nodata, strict=strict, out_dtype=out_dtype)
			else:
				if apply_band < 1 or apply_band > src.count:
					raise ValueError(f"apply-band 超出范围，应在 1..{src.count} 之间 (文件: {in_path})")
				arr = src.read(apply_band)
				res_band = apply_mapping_to_array(arr, mapping=mapping, nodata=nodata, strict=strict, out_dtype=out_dtype)
				full = src.read()
				if full.dtype != out_dtype:
					full = full.astype(out_dtype)
				full[apply_band - 1, :, :] = res_band
				res = full
		write_raster(out_path, res, meta)
		print(f"输出：{out_path}")
		written.append(out_path)
	return written


def parse_args() -> argparse.Namespace:
	parser = argparse.ArgumentParser(description="使用Excel映射表替换GeoTIFF像元值")
	parser.add_argument("--input", help="输入 GeoTIFF 路径")
	parser.add_argument("--input-dir", help="输入目录（批量处理所有 .tif/.tiff）")
	parser.add_argument("--excel", help="Excel 映射表路径")
	parser.add_argument("--output", help="输出 GeoTIFF 路径")
	parser.add_argument("--output-dir", help="输出目录（与 --input-dir 搭配使用）")
	parser.add_argument("--sheet", default="Sheet1", help="Excel 工作表名，默认 Sheet1")
	parser.add_argument("--from-col", dest="from_col", default="from", help="Excel 原值列名，默认 from")
	parser.add_argument("--to-col", dest="to_col", default="to", help="Excel 新值列名，默认 to")
	parser.add_argument("--apply-band", type=int, default=None, help="仅处理的波段序号(1开始)，默认处理全部")
	parser.add_argument("--ignore-missing", action="store_true", help="映射中找不到的值保持不变（默认行为）")
	parser.add_argument("--strict", action="store_true", help="严格模式：出现未映射值直接报错")
	parser.add_argument("--make-sample", action="store_true", help="生成示例TIF与Excel")
	parser.add_argument("--outdir", default="data", help="生成示例文件输出目录，默认 data")
	return parser.parse_args()


def main() -> None:
	args = parse_args()
	if args.make_sample:
		in_tif, in_xlsx = make_sample_data(args.outdir)
		print(f"示例数据已生成：\n  TIFF: {in_tif}\n  Excel: {in_xlsx}")
		return

	# 批量模式
	if args.input_dir and args.output_dir:
		if not args.excel:
			raise SystemExit("批量模式请提供 --excel")
		written = process_directory(
			input_dir=args.input_dir,
			excel_path=args.excel,
			output_dir=args.output_dir,
			sheet=args.sheet,
			from_col=args.from_col,
			to_col=args.to_col,
			apply_band=args.apply_band,
			ignore_missing=args.ignore_missing,
			strict=args.strict,
		)
		print(f"完成：批量输出 {len(written)} 个文件到 {args.output_dir}")
		return

	# 单文件模式
	if args.input and args.output and args.excel:
		process(
			input_tif=args.input,
			excel_path=args.excel,
			output_tif=args.output,
			sheet=args.sheet,
			from_col=args.from_col,
			to_col=args.to_col,
			apply_band=args.apply_band,
			ignore_missing=args.ignore_missing,
			strict=args.strict,
		)
		print(f"完成：已输出 {args.output}")
		return

	raise SystemExit("请提供以下两种用法之一：\n1) 单文件：--input --excel --output\n2) 批量：--input-dir --excel --output-dir 或先使用 --make-sample 生成示例数据")


if __name__ == "__main__":
	main()
